# -*- coding: utf-8 -*-
"""
道路裂缝检测与修复系统 - 基于深度学习的裂缝分割模型
核心功能：使用预训练模型对道路图像进行裂缝检测，输出裂缝概率图

 输入大图，训练完自动拼接
"""

import sys
import os
# 将项目根目录添加到Python搜索路径
# sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import cv2  # OpenCV库，用于图像处理
import numpy as np  # 数值计算库
import torch  # PyTorch深度学习框架
from matplotlib import pyplot as plt

from model.deepcrack_with_uu_Sp_3 import DeepCrack_uu_Sp_3  # 自定义的裂缝分割模型架构
from tqdm import tqdm  # 进度条显示工具
from trainer.trainer import DeepCrackTrainer  # 模型训练器
from config.config import Config as cfg  # 配置文件（包含模型参数、训练参数等）

# 设置CUDA可见设备（指定使用的GPU）
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu_id


def repair(input_path, output_path):
    """
    主处理函数 - 执行裂缝检测全流程

    参数:
        input_path: str - 输入图像目录路径
        output_path: str - 结果保存目录路径
    """

    # ==================== 1. 模型初始化配置 ====================
    # 预训练模型权重路径（注意：路径需根据实际位置修改）
    pretrained_model = r'predicted_model/DuSp3_CRKWH1001_Wight6_2_200_DuSp3_CRKWH1001_Wight6_2_200_epoch(199)_0000200_2025-05-08-13-25-51.pth'

    # 创建输出目录（如果不存在）
    if not os.path.exists(output_path):
        os.makedirs(output_path)  # 递归创建目录

    # ==================== 2. 模型加载与配置 ====================
    # 设置计算设备（优先使用GPU）
    device = torch.device("cuda")

    # 初始化模型架构
    model = DeepCrack_uu_Sp_3(cfg)  # 使用配置文件参数构建模型

    # 多GPU并行计算配置（device_ids指定使用的GPU索引）
    model = torch.nn.DataParallel(model, device_ids=[0]).cuda()

    # 初始化训练器并加载预训练权重
    trainer = DeepCrackTrainer(model).cuda()
    model.load_state_dict(
        trainer.saver.load(pretrained_model, multi_gpu=True),
        strict=True  # 严格匹配权重名称
    )
    model.eval()  # 设置为评估模式（关闭dropout等训练专用层）

    # ==================== 3. 图像处理流水线 ====================
    with torch.no_grad():  # 禁用梯度计算（节省内存）
        # 获取输入目录下所有文件
        files = os.listdir(input_path)
        # 按照文件名排序
        files.sort()
        # print(len(files))
        # 没有文件就退出
        if not files:
            print("输入目录为空，请检查路径！")
            return

        # 获取前pre_num张图像进行检测（默认检测前10张）
        pre_num = 1
        # print(files[:pre_num])
        # 使用进度条遍历处理最后pre_num张图像
        # for file in tqdm(files[-pre_num:], desc="裂缝检测进度"):# 这是从后往前

        for file in tqdm(files[:pre_num], desc="裂缝检测进度"): # 这是从前往后
            # 3.1 文件路径处理
            save_path = os.path.normpath(os.path.join(output_path, file))  # 规范化路径
            img_path = os.path.join(input_path, file)  # 输入图像完整路径

            # 3.2 图像读取与预处理
            image = cv2.imread(img_path)  # 读取BGR格式图像
            w, h, c = image.shape  # 获取图像尺寸（宽度、高度、通道数）

            # 初始化结果存储矩阵（全零张量）
            save_pred = torch.zeros(w, h)


            """
            将大图像分割为stridexstride的小块进行预测（滑动窗口）。
            对每个小块进行归一化并输入模型。
            使用sigmoid函数将输出转换为概率图。
            将预测结果拼接回原图大小。
            """
            numm = 1
            stride = 512
            # 3.3 滑动窗口处理（stride x stride分块预测）
            for i in range(0, w, stride):  # 垂直方向滑动
                for j in range(0, h, stride):  # 水平方向滑动
                    # 提取当前窗口图像
                    image_split = image[i:i + stride, j:j + stride]
                    a, b, c = image_split.shape  # 实际窗口尺寸

                    # 3.4 图像填充（不足stride x stride的部分补零）
                    im = np.zeros([stride, stride, 3], dtype=np.float32)
                    im[0:a, 0:b] = image_split  # 左上角填充有效区域

                    # 3.5 数据标准化与格式转换
                    print(im.shape)
                    img = im.transpose(2, 0, 1)  # HWC -> CHW（PyTorch标准格式）
                    print(img.shape)
                    print(type(img))
                    img = img.astype(np.float32) / 255.0  # 归一化到[0,1]

                    # 转换为PyTorch张量并添加batch维度
                    test_data = torch.from_numpy(img).float().unsqueeze(0).to(device)

                    # 3.6 模型推理
                    # print(type(test_data))
                    # print(test_data.shape)
                    test_pred = model(test_data)  # 前向传播
                    test_pred = torch.sigmoid(test_pred[0].cpu().squeeze())  # 压缩维度并应用sigmoid
                    # print(test_pred)
                    # print(test_pred.shape)

                    # 显示全部小图
                    # plt.figure(figsize=(w / 100, h / 100))
                    # plt.imshow(test_pred)
                    # numm = numm + 1
                    # plt.title('test_pred_'+str(numm))
                    # plt.show()
                    # 3.7 结果拼接
                    if a != stride or b != stride:  # 边缘区块处理
                        save_pred[i:i + a, j:j + b] = test_pred[0:a, 0:b]
                    else:
                        save_pred[i:i + stride, j:j + stride] = test_pred

            # ==================== 4. 结果后处理与保存 ====================
            """
            将预测结果保存为PNG图像。
            对预测结果进行后处理（取0.7次方以增强对比度）。
            将概率图转换为0-255的灰度图并保存。
            """
            # 生成输出文件名（保持原文件名，扩展名改为.png）
            save_name = os.path.join(output_path, file.split('.')[0] + '.png')

            # 概率图增强（gamma校正提升对比度）
            print(save_pred.shape)
            # pree = np.power(save_pred.data.numpy(), 0.7)  # 指数0.7为经验值

            # 转换为0-255范围的灰度图
            save_result = (save_pred.data.numpy() * 255).astype(np.uint8)
            # plt显示 原图尺寸
            plt.figure(figsize=(w/100, h/100))
            plt.imshow(save_result)
            plt.show()

            # 保存结果图像
            cv2.imwrite(save_name, save_result)


if __name__ == '__main__':
    """
    输入: 一组道路图像。
    输出: 对应的裂缝检测结果（灰度图，白色表示裂缝）。
    模型: 基于DeepCrack和DANet的改进模型，用于裂缝检测。
    处理方式: 滑动窗口预测，适用于大尺寸图像。
    """
    # 输入输出路径配置（示例路径，实际使用时需修改）
    input_path = r'../tools/XJU1_jsonToMask/train_img_big'  # 输入图像目录
    output_path = r'test_data/out_big'  # 结果输出目录

    # 执行裂缝检测流程
    repair(input_path, output_path)